#!/usr/bin/env python3
"""
First-Order PBPK Modeling Script

基于simple_test.py流程的一级动力学PBPK建模脚本
该脚本实现简化的一级动力学PBPK模型，专注于基本的一室模型
并生成可视化图表

作者：基于BlackCat@CPPO的PBPK系统
版本：1.0.0
日期：2025-07-24
"""

import os
import sys
import numpy as np
import matplotlib.pyplot as plt
from scipy.integrate import solve_ivp
import pandas as pd
from datetime import datetime

# 添加模块路径
current_dir = os.path.dirname(os.path.abspath(__file__))
if current_dir not in sys.path:
    sys.path.insert(0, current_dir)

class FirstOrderPBPKModel:
    """
    一级动力学PBPK模型
    
    实现简化的PBPK模型，包含：
    - 一级吸收过程
    - 一室分布模型
    - 一级消除过程
    """
    
    def __init__(self, parameters=None):
        """
        初始化一级PBPK模型
        
        参数:
            parameters (dict): 模型参数字典
        """
        self.params = parameters or self.get_default_parameters()
        self.time_points = None
        self.results = {}
        
    def get_default_parameters(self):
        """获取默认模型参数"""
        return {
            'dose': 100.0,           # 给药剂量 (mg)
            'ka': 1.0,              # 吸收速率常数 (1/h)
            'ke': 0.2,              # 消除速率常数 (1/h)
            'Vd': 50.0,             # 分布容积 (L)
            'tlag': 0.1,            # 吸收滞后时间 (h)
            'F': 0.8,               # 生物利用度
            'body_weight': 70.0,    # 体重 (kg)
            'age': 30.0,            # 年龄 (岁)
            'gender': 'male'        # 性别
        }
    
    def first_order_absorption(self, t, ka=None, dose=None, tlag=None, F=None):
        """
        计算一级吸收过程
        
        参数:
            t: 时间数组
            ka: 吸收速率常数
            dose: 剂量
            tlag: 滞后时间
            F: 生物利用度
            
        返回:
            dict: 吸收过程结果
        """
        ka = ka or self.params['ka']
        dose = dose or self.params['dose']
        tlag = tlag or self.params['tlag']
        F = F or self.params['F']
        
        # 考虑滞后时间
        t_eff = np.maximum(t - tlag, 0)
        
        # 一级吸收公式
        absorbed_amount = F * dose * (1 - np.exp(-ka * t_eff))
        absorption_rate = F * dose * ka * np.exp(-ka * t_eff)
        
        return {
            'absorbed_amount': absorbed_amount,
            'absorption_rate': absorption_rate,
            'fraction_absorbed': absorbed_amount / (F * dose)
        }
    
    def one_compartment_model(self, t, dose=None, ka=None, ke=None, Vd=None, F=None, tlag=None):
        """
        一室模型PBPK计算
        
        参数:
            t: 时间数组
            dose: 剂量
            ka: 吸收速率常数
            ke: 消除速率常数
            Vd: 分布容积
            F: 生物利用度
            tlag: 滞后时间
            
        返回:
            dict: 一室模型结果
        """
        dose = dose or self.params['dose']
        ka = ka or self.params['ka']
        ke = ke or self.params['ke']
        Vd = Vd or self.params['Vd']
        F = F or self.params['F']
        tlag = tlag or self.params['tlag']
        
        # 考虑滞后时间
        t_eff = np.maximum(t - tlag, 0)
        
        # 计算血浆浓度
        # C(t) = (F*D*ka)/(Vd*(ka-ke)) * (e^(-ke*t) - e^(-ka*t))
        if abs(ka - ke) < 1e-6:  # 防止除零错误
            concentration = (F * dose * ka * t_eff * np.exp(-ka * t_eff)) / Vd
        else:
            concentration = (F * dose * ka) / (Vd * (ka - ke)) * (
                np.exp(-ke * t_eff) - np.exp(-ka * t_eff)
            )
        
        # 计算AUC
        auc = (F * dose) / (ke * Vd)
        
        # 计算消除半衰期
        t_half = 0.693 / ke
        
        # 计算清除率
        clearance = ke * Vd
        
        return {
            'concentration': concentration,
            'auc': auc,
            't_half': t_half,
            'clearance': clearance,
            'cmax': np.max(concentration),
            'tmax': t[np.argmax(concentration)]
        }
    
    def run_simulation(self, simulation_time=24, time_step=0.1):
        """
        运行PBPK模拟
        
        参数:
            simulation_time: 模拟时间 (小时)
            time_step: 时间步长 (小时)
            
        返回:
            dict: 模拟结果
        """
        # 生成时间序列
        self.time_points = np.arange(0, simulation_time + time_step, time_step)
        
        # 计算吸收过程
        absorption_result = self.first_order_absorption(self.time_points)
        
        # 计算一室模型
        pk_result = self.one_compartment_model(self.time_points)
        
        # 整合结果
        self.results = {
            'time': self.time_points,
            'absorption': absorption_result,
            'pharmacokinetics': pk_result,
            'parameters': self.params
        }
        
        return self.results
    
    def generate_plots(self, output_dir='../output'):
        """
        生成可视化图表
        
        参数:
            output_dir: 输出目录
        """
        if not self.results:
            print("Please run simulation first")
            return
        
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)
        
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        
        # Create plots
        fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2, 2, figsize=(15, 12))
        
        time = self.results['time']
        concentration = self.results['pharmacokinetics']['concentration']
        absorbed = self.results['absorption']['absorbed_amount']
        absorption_rate = self.results['absorption']['absorption_rate']
        
        # 1. Concentration-time curve
        ax1.plot(time, concentration, 'b-', linewidth=2, label='Plasma Concentration')
        ax1.set_xlabel('Time (hours)')
        ax1.set_ylabel('Concentration (mg/L)')
        ax1.set_title('Plasma Concentration-Time Curve')
        ax1.grid(True, alpha=0.3)
        ax1.legend()
        
        # 2. Absorption process
        ax2.plot(time, absorbed, 'g-', linewidth=2, label='Cumulative Absorption')
        ax2_twin = ax2.twinx()
        ax2_twin.plot(time, absorption_rate, 'r--', linewidth=2, label='Absorption Rate')
        ax2.set_xlabel('Time (hours)')
        ax2.set_ylabel('Absorbed Amount (mg)', color='g')
        ax2_twin.set_ylabel('Absorption Rate (mg/h)', color='r')
        ax2.set_title('Drug Absorption Process')
        ax2.grid(True, alpha=0.3)
        
        # 3. Semi-log concentration-time curve
        ax3.semilogy(time, concentration, 'b-', linewidth=2)
        ax3.set_xlabel('Time (hours)')
        ax3.set_ylabel('Concentration (mg/L)')
        ax3.set_title('Semi-log Concentration-Time Curve')
        ax3.grid(True, alpha=0.3)
        
        # 4. Drug disposition process
        # Calculate elimination phase
        ke = self.params['ke']
        elimination_rate = ke * concentration * self.params['Vd']
        
        ax4.plot(time, absorption_rate, 'g-', linewidth=2, label='Absorption Rate')
        ax4.plot(time, elimination_rate, 'r-', linewidth=2, label='Elimination Rate')
        ax4.set_xlabel('Time (hours)')
        ax4.set_ylabel('Rate (mg/h)')
        ax4.set_title('Drug Disposition Process')
        ax4.legend()
        ax4.grid(True, alpha=0.3)
        
        plt.tight_layout()
        
        # 保存图表
        filename = os.path.join(output_dir, f'first_order_pbpk_{timestamp}.png')
        plt.savefig(filename, dpi=300, bbox_inches='tight')
        plt.show()
        
        print(f"Chart saved to: {filename}")
        
        return filename
    
    def export_results(self, output_dir='../output'):
        """
        导出模拟结果
        
        参数:
            output_dir: 输出目录
        """
        if not self.results:
            print("Please run simulation first")
            return
        
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)
        
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        
        # 创建DataFrame
        df = pd.DataFrame({
            'time': self.results['time'],
            'concentration': self.results['pharmacokinetics']['concentration'],
            'absorbed_amount': self.results['absorption']['absorbed_amount'],
            'absorption_rate': self.results['absorption']['absorption_rate']
        })
        
        # 导出CSV
        csv_file = os.path.join(output_dir, f'first_order_results_{timestamp}.csv')
        df.to_csv(csv_file, index=False)
        
        # 导出参数
        params_file = os.path.join(output_dir, f'parameters_{timestamp}.json')
        import json
        with open(params_file, 'w', encoding='utf-8') as f:
            json.dump(self.params, f, indent=2, ensure_ascii=False)
        
        print(f"Results exported to: {csv_file}")
        print(f"Parameters exported to: {params_file}")
        
        return csv_file, params_file
    
    def run_parameter_sensitivity(self, parameter_name, parameter_range, simulation_time=24):
        """
        运行参数敏感性分析
        
        参数:
            parameter_name: 参数名称
            parameter_range: 参数范围列表
            simulation_time: 模拟时间
        """
        results = []
        
        for param_value in parameter_range:
            # 临时修改参数
            original_value = self.params[parameter_name]
            self.params[parameter_name] = param_value
            
            # 运行模拟
            sim_result = self.run_simulation(simulation_time)
            
            # 记录结果
            results.append({
                'parameter_value': param_value,
                'cmax': sim_result['pharmacokinetics']['cmax'],
                'auc': sim_result['pharmacokinetics']['auc'],
                't_half': sim_result['pharmacokinetics']['t_half']
            })
            
            # 恢复参数
            self.params[parameter_name] = original_value
        
        return results

def main():
    """主函数"""
    print("=== First-Order PBPK Modeling System ===")
    
    # 创建模型实例
    model = FirstOrderPBPKModel()
    
    # 设置自定义参数（可选）
    custom_params = {
        'dose': 200.0,      # 200 mg
        'ka': 1.5,          # 1.5 1/h
        'ke': 0.15,         # 0.15 1/h
        'Vd': 40.0,         # 40 L
        'tlag': 0.2,        # 0.2 h
        'F': 0.85           # 85% 生物利用度
    }
    model.params.update(custom_params)
    
    # 运行模拟
    print("运行PBPK模拟...")
    results = model.run_simulation(simulation_time=24, time_step=0.1)
    
    # 显示结果摘要
    print("\n=== 模拟结果摘要 ===")
    pk = results['pharmacokinetics']
    print(f"最大浓度 (Cmax): {pk['cmax']:.2f} mg/L")
    print(f"达峰时间 (Tmax): {pk['tmax']:.2f} h")
    print(f"消除半衰期 (t1/2): {pk['t_half']:.2f} h")
    print(f"AUC: {pk['auc']:.2f} mg·h/L")
    print(f"清除率: {pk['clearance']:.2f} L/h")
    
    # 生成图表
    print("\n生成可视化图表...")
    model.generate_plots()
    
    # 导出结果
    print("\n导出结果...")
    model.export_results()
    
    print("\n=== 模拟完成 ===")
    
    # 运行参数敏感性分析（可选）
    print("\n运行ka参数敏感性分析...")
    ka_range = [0.5, 1.0, 1.5, 2.0, 2.5]
    sensitivity_results = model.run_parameter_sensitivity('ka', ka_range)
    
    print("\nka敏感性分析结果:")
    for result in sensitivity_results:
        print(f"ka={result['parameter_value']}: Cmax={result['cmax']:.2f}, AUC={result['auc']:.2f}")

if __name__ == "__main__":
    main()